import os
import numpy as np
from scipy.fft import fft
from scipy.spatial.distance import cosine
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.size'] = 10

def read_data_from_folder(folder_path):
    data_list = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith('.txt'):
            file_path = os.path.join(folder_path, file_name)
            data = np.loadtxt(file_path)
            data_list.append(data)
    return np.array(data_list)

def calculate_time_domain_features(data):
    mean_val = np.mean(data, axis=1)
    std_val = np.std(data, axis=1)
    max_val = np.max(data, axis=1)
    return mean_val, std_val, max_val

def calculate_frequency_domain_features(data):
    fft_data = np.abs(fft(data, axis=1))
    n = data.shape[1]
    fft_data = fft_data[:, :n // 2]
    mean_freq = np.mean(fft_data, axis=1)
    std_freq = np.std(fft_data, axis=1)
    max_freq = np.max(fft_data, axis=1)

    return mean_freq, std_freq, max_freq

def calculate_cosine_similarity_to_mean(data, class_mean):
    similarities = []
    for sample in data:
        cos_distance = cosine(sample, class_mean)
        cos_similarity = 1 - cos_distance
        similarities.append(cos_similarity)
    return np.array(similarities)

def train_model(class0_path, class1_path):
    class0_data = read_data_from_folder(class0_path)
    class1_data = read_data_from_folder(class1_path)
    print(f"训练集 - Class 0 样本数: {class0_data.shape[0]}")
    print(f"训练集 - Class 1 样本数: {class1_data.shape[0]}")
    mean0, std0, max0 = calculate_time_domain_features(class0_data)
    mean1, std1, max1 = calculate_time_domain_features(class1_data)
    mean_freq0, std_freq0, max_freq0 = calculate_frequency_domain_features(class0_data)
    mean_freq1, std_freq1, max_freq1 = calculate_frequency_domain_features(class1_data)
    class0_mean = np.mean(class0_data, axis=0)
    class1_mean = np.mean(class1_data, axis=0)
    cos_sim0 = calculate_cosine_similarity_to_mean(class0_data, class0_mean)
    cos_sim1 = calculate_cosine_similarity_to_mean(class1_data, class1_mean)
    class0_fft = np.abs(fft(class0_data, axis=1))
    class1_fft = np.abs(fft(class1_data, axis=1))
    class0_fft_mean = np.mean(class0_fft, axis=0)
    class1_fft_mean = np.mean(class1_fft, axis=0)
    cos_sim_freq0 = calculate_cosine_similarity_to_mean(class0_fft, class0_fft_mean)
    cos_sim_freq1 = calculate_cosine_similarity_to_mean(class1_fft, class1_fft_mean)
    mean0_mean, mean1_mean = np.mean(mean0), np.mean(mean1)
    std0_mean, std1_mean = np.mean(std0), np.mean(std1)
    max0_mean, max1_mean = np.mean(max0), np.mean(max1)
    cos_sim0_mean, cos_sim1_mean = np.mean(cos_sim0), np.mean(cos_sim1)
    mean_freq0_mean, mean_freq1_mean = np.mean(mean_freq0), np.mean(mean_freq1)
    std_freq0_mean, std_freq1_mean = np.mean(std_freq0), np.mean(std_freq1)
    max_freq0_mean, max_freq1_mean = np.mean(max_freq0), np.mean(max_freq1)
    cos_sim_freq0_mean, cos_sim_freq1_mean = np.mean(cos_sim_freq0), np.mean(cos_sim_freq1)
    thresholds = {
        'mean_threshold': (mean0_mean + mean1_mean) / 2,
        'std_threshold': (std0_mean + std1_mean) / 2,
        'max_threshold': (max0_mean + max1_mean) / 2,
        'cos_sim_threshold': (cos_sim0_mean + cos_sim1_mean) / 2,
        'mean_freq_threshold': (mean_freq0_mean + mean_freq1_mean) / 2,
        'std_freq_threshold': (std_freq0_mean + std_freq1_mean) / 2,
        'max_freq_threshold': (max_freq0_mean + max_freq1_mean) / 2,
        'cos_sim_freq_threshold': (cos_sim_freq0_mean + cos_sim_freq1_mean) / 2
    }
    # 如果class0均值 < class1均值，则样本特征值小于阈值时判为class0，否则判为class1
    directions = {
        'mean_direction': 0 if mean0_mean < mean1_mean else 1,
        'std_direction': 0 if std0_mean < std1_mean else 1,
        'max_direction': 0 if max0_mean < max1_mean else 1,
        'cos_sim_direction': 0 if cos_sim0_mean < cos_sim1_mean else 1,
        'mean_freq_direction': 0 if mean_freq0_mean < mean_freq1_mean else 1,
        'std_freq_direction': 0 if std_freq0_mean < std_freq1_mean else 1,
        'max_freq_direction': 0 if max_freq0_mean < max_freq1_mean else 1,
        'cos_sim_freq_direction': 0 if cos_sim_freq0_mean < cos_sim_freq1_mean else 1
    }
    class_means = {
        'class0_mean': class0_mean,
        'class1_mean': class1_mean,
        'class0_fft_mean': class0_fft_mean,
        'class1_fft_mean': class1_fft_mean
    }
    return thresholds, directions, class_means

def calculate_sample_features(sample, class_means):
    mean_val = np.mean(sample)
    std_val = np.std(sample)
    max_val = np.max(sample)
    cos_sim_time_0 = 1 - cosine(sample, class_means['class0_mean'])
    cos_sim_time_1 = 1 - cosine(sample, class_means['class1_mean'])
    cos_sim_time = cos_sim_time_0
    fft_sample = np.abs(fft(sample))
    n = len(sample)
    fft_sample = fft_sample[:n // 2]
    mean_freq = np.mean(fft_sample)
    std_freq = np.std(fft_sample)
    max_freq = np.max(fft_sample)
    cos_sim_freq_0 = 1 - cosine(fft_sample, class_means['class0_fft_mean'][:n // 2])
    cos_sim_freq_1 = 1 - cosine(fft_sample, class_means['class1_fft_mean'][:n // 2])
    cos_sim_freq = cos_sim_freq_0

    features = {
        'mean': mean_val,
        'std': std_val,
        'max': max_val,
        'cos_sim': cos_sim_time,
        'mean_freq': mean_freq,
        'std_freq': std_freq,
        'max_freq': max_freq,
        'cos_sim_freq': cos_sim_freq
    }

    return features

def testmodel(class0_path, class1_path, thresholds, directions, class_means):
    class0_test = read_data_from_folder(class0_path)
    class1_test = read_data_from_folder(class1_path)
    print(f"测试集 - Class 0 样本数: {class0_test.shape[0]}")
    print(f"测试集 - Class 1 样本数: {class1_test.shape[0]}")

    test_data = np.concatenate([class0_test, class1_test])
    test_labels = np.concatenate([
        np.zeros(class0_test.shape[0]),
        np.ones(class1_test.shape[0])
    ])

    predictions = []
    votes = []

    for sample in test_data:
        features = calculate_sample_features(sample, class_means)
        feature_votes = []
        if directions['mean_direction'] == 0:
            vote = 0 if features['mean'] < thresholds['mean_threshold'] else 1
        else:
            vote = 1 if features['mean'] < thresholds['mean_threshold'] else 0
        feature_votes.append(vote)
        if directions['std_direction'] == 0:
            vote = 0 if features['std'] < thresholds['std_threshold'] else 1
        else:
            vote = 1 if features['std'] < thresholds['std_threshold'] else 0
        feature_votes.append(vote)
        if directions['max_direction'] == 0:
            vote = 0 if features['max'] < thresholds['max_threshold'] else 1
        else:
            vote = 1 if features['max'] < thresholds['max_threshold'] else 0
        feature_votes.append(vote)
        if directions['cos_sim_direction'] == 0:
            vote = 0 if features['cos_sim'] < thresholds['cos_sim_threshold'] else 1
        else:
            vote = 1 if features['cos_sim'] < thresholds['cos_sim_threshold'] else 0
        feature_votes.append(vote)
        if directions['mean_freq_direction'] == 0:
            vote = 0 if features['mean_freq'] < thresholds['mean_freq_threshold'] else 1
        else:
            vote = 1 if features['mean_freq'] < thresholds['mean_freq_threshold'] else 0
        feature_votes.append(vote)
        if directions['std_freq_direction'] == 0:
            vote = 0 if features['std_freq'] < thresholds['std_freq_threshold'] else 1
        else:
            vote = 1 if features['std_freq'] < thresholds['std_freq_threshold'] else 0
        feature_votes.append(vote)
        if directions['max_freq_direction'] == 0:
            vote = 0 if features['max_freq'] < thresholds['max_freq_threshold'] else 1
        else:
            vote = 1 if features['max_freq'] < thresholds['max_freq_threshold'] else 0
        feature_votes.append(vote)
        if directions['cos_sim_freq_direction'] == 0:
            vote = 0 if features['cos_sim_freq'] < thresholds['cos_sim_freq_threshold'] else 1
        else:
            vote = 1 if features['cos_sim_freq'] < thresholds['cos_sim_freq_threshold'] else 0
        feature_votes.append(vote)
        votes.append(feature_votes)
        final_vote = 0 if sum(feature_votes) < 4 else 1
        predictions.append(final_vote)
    predictions = np.array(predictions)
    votes = np.array(votes)

    return test_labels, predictions, votes

def main():
    class0_train_path = r"D:\A_90762\cui_python\data\train\class0"
    class1_train_path = r"D:\A_90762\cui_python\data\train\class1"
    class0_test_path = r"D:\A_90762\cui_python\data\test\class0"
    class1_test_path = r"D:\A_90762\cui_python\data\test\class1"

    print("正在训练模型并计算阈值...")
    thresholds, directions, class_means = train_model(class0_train_path, class1_train_path)

    print("\n训练得到的阈值:")
    for key, value in thresholds.items():
        print(f"{key}: {value:.6f}")

    print("\n正在测试模型...")
    test_labels, predictions, votes = testmodel(
        class0_test_path, class1_test_path, thresholds, directions, class_means
    )

    accuracy = accuracy_score(test_labels, predictions)
    print(f"\n整体准确率: {accuracy:.4f}")

    cm = confusion_matrix(test_labels, predictions)
    print("\n混淆矩阵:")
    print(cm)

    print("\n分类报告:")
    print(classification_report(test_labels, predictions, target_names=['Class 0', 'Class 1']))

    print("\n各特征单独分类准确率:")
    feature_names = [
        '时域均值', '时域标准差', '时域最大值', '时域余弦相似度',
        '频域均值', '频域标准差', '频域最大值', '频域余弦相似度'
    ]

    for i, name in enumerate(feature_names):
        feature_accuracy = accuracy_score(test_labels, votes[:, i])
        print(f"{name}: {feature_accuracy:.4f}")

    plt.figure(figsize=(10, 6))
    accuracies = [accuracy_score(test_labels, votes[:, i]) for i in range(8)]
    plt.bar(range(8), accuracies)
    plt.axhline(y=accuracy, color='r', linestyle='--', label='整体准确率')
    plt.xticks(range(8), feature_names, rotation=45)
    plt.ylabel('准确率')
    plt.title('各特征单独分类准确率')
    plt.legend()
    plt.tight_layout()
    plt.savefig('feature_accuracies.png', dpi=300)
    plt.show()

if __name__ == "__main__":
    main()

